# -*- coding: utf-8 -*-

import numpy as np
  
L1 = 0.544/2 # m
OA = 0.717/2 # m


W1 = np.array([OA, L1, 0])
W2 = np.array([OA, -L1, 0])
W3 = np.array([-OA, L1, 0])
W4 = np.array([-OA, -L1, 0])

W_list = [W1, W2, W3, W4]

wheel_num = len(W_list)


wz_unit_up = np.array([0, 0, 1])
wz_unit_do = np.array([0, 0, -1])

O = np.array([0, 0, 0]) # 坐标原点

## ================================== 公用函数

def get_theta(vec):
    # 输入一个 单位化的向量，输出其与 x 轴正方向的夹角。
    # 返回的角度是度数 

    # assert x^2 + y^2 = 1 

    x = vec[0]
    y = vec[1]

    theta = np.arcsin(y) # y/1   rad
    if y > 0 and x > 0:
        theta = np.rad2deg(theta) # 0 - 90
    if y > 0 and x < 0:
        theta = 180 - np.rad2deg(theta)
    if y < 0 and x > 0:
        theta = np.rad2deg(theta) # -0 - -90
    if y < 0 and x < 0:
        theta = -( 180-abs(np.rad2deg(theta)) )
        
    if y == 0:
        if x > 0:
            theta = 0
        if x < 0:
            theta = 180

    if x == 0:
        if y > 0:
            theta = 90
        if y < 0:
            theta = -90
    
    return theta

# 四个轮子每个轮子都有自己的角度范围
# w1 0-30  90-180   0- -90   -150- -180
# w2 0-90  150-180   0- -30  -90  - -180  
# w3 0-90  150-180   0- -30  -90  - -180  
# w4 0-30  90-180   0- -90   -150- -180
# 对角是一样的， w1,w4,   w2,w3  是一样的 

def check_w1(theta1):
    # 传进来的是角度 
    if -90 <= theta1 <= 30 or 90<= theta1 <= 180 or -180 <= theta1 <= -150:
        return True 
    else:
        return False 

def check_w2(theta2):
    if -30 <= theta2 <= 90 or 150 <= theta2 <= 180 or -180 <= theta2 <= -90:
        return True 
    else:
        return False 

def check_w3(theta3):
    return check_w2(theta3)

def check_w4(theta4):
    return check_w1(theta4) 

def check_all_theta(theta_list):

    return [check_w1(theta_list[0]) ,
            check_w2(theta_list[1]) ,
            check_w3(theta_list[2]) ,
            check_w4(theta_list[3])]


def calculate_diagonal_deg(theta):
    # diag_deg = calculate_diagonal_deg(theta)

    # 求一个角的对角

    # 输入的是 0-180， 0- -180 的角度，
    # 比如 3度，计算对角得到 -177 度

    deg = 180.0 - abs(theta)
    if theta > 0:
        return -deg
    if theta < 0:
        return deg
    if theta == 0:
        return 180


def adjust_v_deg(v_list, theta_list):
    # new_v_list, new_deg_list = adjust_v_deg(v_list, theta_list)
    new_v_list = []
    new_deg_list = []
    
    n = 4
    deg_thresh = [ [-90, 30], [-30, 90], [-30, 90], [-90, 30] ]
    
    for i in range(n):
        v = v_list[i]
        deg = theta_list[i]
        
        deg_min = deg_thresh[i][0]
        deg_max = deg_thresh[i][1]
        
        if deg_min <= deg <= deg_max:
            new_v_list.append(v)
            new_deg_list.append(deg)
        else:
            print('adjust_v_deg:', i, ' deg:', deg)
            new_v = -v
            new_deg = calculate_diagonal_deg(deg)

            new_v_list.append(new_v)
            new_deg_list.append(new_deg)

    return new_v_list, new_deg_list
        
## ==================================

def get_vt(v, w):
    # get v_list and theta_list 
    # v_list, theta_list = get_vt(v, w)
    print("input v,w:", v, w)

    # get v_list and theta_list
    if v == 0 and w == 0:
        print('0, stop mode')
        v_list = [0]*wheel_num
        theta_list = [0]*wheel_num
    
    if v > 0 and w == 0:
        print('1, forward mode')
        v_list = [v]*wheel_num
        theta_list = [0]*wheel_num



    if v < 0 and w == 0:
        print('2, backward mode')
        v_list = [v]*wheel_num
        theta_list = [0]*wheel_num


    if v == 0 and w > 0:
        print('3, turn left inplace mode') # 原地左转
        
        OW = [W-O for W in W_list]
        R = [np.linalg.norm(ow) for ow in OW] 
        WT = []
        # WT = np.cross(OW, wz_do)
        for i in range(wheel_num):
            wt = np.cross(OW[i], W_list[i]+wz_unit_do ) 
            WT.append(wt[:2]) # remove z 
            
        WT = [wt/np.linalg.norm(wt) for wt in WT] # 单位化
    
        v_list = [w*r for r in R]
        theta_list = [get_theta(wt) for wt in WT]
        # 和 x 轴正方向的夹角是正确的

    if v == 0 and w < 0:
        print('4, turn right inplace mode')

        OW = [W-O for W in W_list]
        R = [np.linalg.norm(ow) for ow in OW] 
        WT = []
        # WT = np.cross(OW, wz_do)
        for i in range(wheel_num):
            wt = np.cross(OW[i], W_list[i]+wz_unit_up ) 
            WT.append(wt[:2]) # remove z 
            
        WT = [wt/np.linalg.norm(wt) for wt in WT] # 单位化

        v_list = [abs(w)*r for r in R]
        theta_list = [get_theta(wt) for wt in WT]
        

    if abs(v) > 0 and abs(w) > 0: # 5
        print('5, abs(v) > 0 and abs(w) > 0 mode')
        # 参考 t1.py 
        r = abs(1.0*v/w) # m
        print('r: ', round(r, 2), ' m')
        
        # if abs(r) == L1:
        #     r = abs(L1+0.01)

        ## 假设左转的时候 旋转中心是 O1
        if w > 0:
            sign = 1 
        if w < 0:
            sign = -1 
            
        O1 = np.array([0, sign*r, 0])

        WO1 = [O1-W for W in W_list]
        R = [np.linalg.norm(wo1) for wo1 in WO1]
        WT = []
        for i in range(wheel_num):
            wi_o1 = WO1[i]

            if v*w > 0:
                wz = W_list[i] + wz_unit_up
            if v*w < 0:
                wz = W_list[i] + wz_unit_do

            wi_ti = np.cross(wi_o1, wz) # 叉乘
            
            
            WT.append(wi_ti[:2]/np.linalg.norm(wi_ti[:2]))
            

        v_list = [abs(w*r) for r in R] # m/s
        theta_list = [get_theta(wt) for wt in WT]

    return v_list, theta_list 

def gen_vw_list(v, w):
#    vl, dl = gen_vw_list(v, w)
    assert abs(v/w) > 0.90, "r < 0.893"
    v_list, theta_list = get_vt(v, w) # ground truth
    
    if check_all_theta(theta_list) == [True, True, True, True]:
        new_v_list, new_deg_list = adjust_v_deg(v_list, theta_list)
        
        new_v_list = [round(x, 3) for x in new_v_list]
        new_deg_list = [round(x, 3) for x in new_deg_list]
        
        return new_v_list, new_deg_list
    else:
        return [0]*4, [0]*4 
    